from osgeo import gdal
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.colors import Normalize


def read_raster(raster_path,band):
    dataset = gdal.Open(raster_path)

    if dataset is None:
        raise Exception(f"Unable to open raster file: {raster_path}")

    # 获取地理转换参数
    geo_trans = dataset.GetGeoTransform()

    width = dataset.RasterXSize
    height = dataset.RasterYSize
    num_bands = dataset.RasterCount

    #获取投影
    proj = dataset.GetProjection()

    #获取影像某一个波段
    band_data=dataset.GetRasterBand(band)

    # 以numpy数组方式读入栅格影像
    band_array=band_data.ReadAsArray()

    return geo_trans,proj, num_bands, band_array



def plot_raster(raster_data,color_composite=True):
    """
    绘制栅格数据
    """
    num_bands = len(raster_data)

    if color_composite and num_bands >= 3:
        # 创建彩色合成影像
        rgb_image = np.stack((raster_data[2], raster_data[1], raster_data[0]), axis=-1)
        plt.imshow(rgb_image)
    else:
        # 创建灰度图
        plt.imshow(raster_data[0], cmap='gray')

    plt.show()



def histogram_stretch(raster_data):
    """
    对直方图进行2%线性拉伸
    """
    stretched_data = []

    for band_data in raster_data:
        # Calculate 2% and 98% percentile values
        p2 = np.percentile(band_data, 2)
        p98 = np.percentile(band_data, 98)

        # Linear stretch to range [0, 255]
        stretched_band = np.clip(255 * (band_data - p2) / (p98 - p2), 0, 255)
        stretched_data.append(stretched_band)

    return stretched_data




def write_raster(output_path,band_array,geo_trans,proj):
    driver = gdal.GetDriverByName("GTiff")  ## 创建.tif文件驱动
    outdset = driver.Create(output_path, xsize=band_array.shape[1], \
                                    ysize=band_array.shape[0], bands=1, eType=gdal.GDT_Int16)  ### 创建空的.tif数据
    outdset.SetGeoTransform(geo_trans)    ### 设置地理转换参数
    outdset.SetProjection(proj)           ### 设置投影
    outband = outdset.GetRasterBand(1)    ### 获取波段1
    outband.WriteArray(band_array)        ### 将np.array()数组写入波段1
    outband.SetNoDataValue(np.nan)        ### 设置特定值为无数据，具有压缩功能
    outdset = None    ## 关闭.tif文件驱动
    print("done")



def read_multiband(raster_path, band_indices=(1, 2, 3),stretch=True):
    dataset = gdal.Open(raster_path)

    if dataset is None:
        raise Exception(f"Unable to open raster file: {raster_path}")

    # 获取地理转换参数
    geo_trans = dataset.GetGeoTransform()

    width = dataset.RasterXSize
    height = dataset.RasterYSize
    num_bands = dataset.RasterCount

    # 获取投影
    proj = dataset.GetProjection()

    # 以numpy数组方式读入栅格影像
    raster_data = dataset.ReadAsArray()

    # 选择指定的波段
    selected_bands = [raster_data[band_idx - 1] for band_idx in band_indices]

    # 彩色合成为RGB影像
    rgb_image = np.stack(selected_bands, axis=-1)

        # 对图像进行拉伸显示
    if stretch:
        min_val, max_val = np.percentile(rgb_image,2), np.percentile(rgb_image,98)
        norm = Normalize(vmin=min_val, vmax=max_val)
        stretched_rgb = norm(rgb_image)

        # 显示拉伸后的影像
        plt.imshow(stretched_rgb)
    else:
        # 显示原始影像
        plt.imshow(rgb_image)
        
    plt.title("RGB Composite")
    plt.show()

    return geo_trans, proj, num_bands, rgb_image




